import chromadb

from app.embedding import get_embedding
from app.vectorDB.chunk import chunking

paragraphs = chunking("../../WitSpider/数据集/OrChiD.json")

class VectorDBConnector:
    def __init__(self, connection_name, embedding_function):
        '''初始化方法'''
        # 持久化存储
        chroma_client = chromadb.PersistentClient(path="./chroma")  # path的值即存储路径的值可以自行指定
        # 创建一个collection
        self.collection = chroma_client.get_or_create_collection(name=connection_name)
        self.embedding_function = embedding_function

    def addDE(self, doc):
        '''添加原始文档与对应的向量'''
        self.collection.add(
            embeddings=self.embedding_function(doc),  # 文档对应的向量（使用自己绑定的向量模型，不设置也可，Chroma有内置向量模型）
            documents=doc,  # 原始文档
            ids=[f"id{i}" for i in range(len(doc))]  # 文档索引id
        )

    def search(self, query, top_n):  # query为查询问题/内容，top_n为查询返回最相似的段落个数
        '''检索方法'''
        results = self.collection.query(
            query_embeddings=self.embedding_function(query),  # 问题文本向量化
            n_results=top_n  # 返回个数
        )
        return results


# 创建向量数据库对象，名字为test，向量化函数为之前写好的m3e-base模型的向量化函数
vectorDB = VectorDBConnector("test", get_embedding)
# 添加刚刚已经分块完毕的 paragraphs 文档
vectorDB.addDE(paragraphs[:500])

test_query = "爱需要克制吗"
results = vectorDB.search(test_query, 2)

for para in results['documents'][0]:
    print(para+"\n")